Amazon Aurora DSQLのカスタムデータベースロールを使ってLambdaから接続してみた #AWSreInvent

import * as cdk from 'aws-cdk-lib';
import { Construct } from 'constructs';
import * as iam from 'aws-cdk-lib/aws-iam';
import * as lambda from 'aws-cdk-lib/aws-lambda';
import * as nodejs from 'aws-cdk-lib/aws-lambda-nodejs';

export class DsqlLambdaStack extends cdk.Stack {
  constructor(scope: Construct, id: string, props?: cdk.StackProps) {
    super(scope, id, props);

    const dsqlLambda = new nodejs.NodejsFunction(this, 'DsqlLambda', {
      entry: 'src/index.ts',
      runtime: lambda.Runtime.NODEJS_LATEST,
      environment: {
        DB_USER: 'member',
        DATABASE: 'postgres',
        HOST_NAME: process.env.CLUSTER_ENDPOINT!, // HAK
        REGION: 'us-east-1',
        SCHEMA: 'testschema',
      },
      initialPolicy: [
        new iam.PolicyStatement({
          actions: ['dsql:DbConnect'],
          resources: [
            // See: https://docs.aws.amazon.com/aurora-dsql/latest/userguide/using-iam-condition-keys.html
            this.formatArn({
              service: 'dsql',
              region: 'us-east-1',
              resource: 'cluster',
              resourceName: process.env.CLUSTER_ID!,
            }),
          ],
        }),
      ],
    });

    // Create the role first and then register it with the DSQL role
    new cdk.CfnOutput(this, 'LambdaRoleArn', {
      value: dsqlLambda.role!.roleArn,
    });
  }
}

import { DsqlSigner } from '@aws-sdk/dsql-signer';
import { DataSource } from 'typeorm';

const username = process.env.DB_USER;
const database = process.env.DATABASE;
const hostname = process.env.HOST_NAME!;
const region = process.env.REGION!;
const schema = process.env.SCHEMA!;

const generateToken = async (hostname: string, region: string) => {
  const signer = new DsqlSigner({
    hostname: hostname,
    region,
  });
  try {
    // NOTE: According to the official doc, it requires arguments
    //       but it does not work unless executed without arguments
    const token = await signer.getDbConnectAuthToken();
    return token;
  } catch (error) {
    console.error('Failed to generate token: ', error);
    throw error;
  }
};

let dataSource: DataSource;
const getDataSource = async (): Promise<DataSource> => {
  if (dataSource?.isInitialized) {
    return dataSource;
  }

  const password = await generateToken(hostname, region);
  dataSource = new DataSource({
    type: 'postgres',
    host: hostname,
    port: 5432,
    username,
    password,
    database,
    schema,
    ssl: true,
  });
  await dataSource.initialize();
  return dataSource;
};

export const handler = async () => {
  const ds = await getDataSource();
  const result = await ds.query('SELECT * FROM testschema.users');
  console.log('SQL result:', result);
};

% export CLUSTER_ID='xxxxxxxxxx'
% export CLUSTER_ENDPOINT='xxxxxxxxxx.dsql.us-east-1.on.aws'
% npx cdk deploy

Outputs:
DsqlLambdaStack.LambdaRoleArn = arn:aws:iam::111122223333:role/DsqlLambdaStack-DsqlLambdaServiceRoleXXX

% aws dsql generate-db-connect-admin-auth-token \
  --region us-east-1 \
  --expires 3600 \
  --hostname xxxxxxxxxx.dsql.us-east-1.on.aws

% psql --host=xxxxxxxxxx.dsql.us-east-1.on.aws \
  --port=5432 \
  --dbname=postgres \
  --username=admin

postgres=> 
CREATE SCHEMA testschema
CREATE TABLE users(user_id int, name text);
INSERT INTO users (user_id, name) VALUES 
    (1, 'John Smith'),
    (2, 'Maria Garcia'),
    (3, 'Yuki Tanaka'),
    (4, 'David Wilson'),
    (5, 'Sarah Johnson');

CREATE ROLE member WITH LOGIN;
AWS IAM GRANT member TO 'arn:aws:iam::111122223333:role/DsqlLambdaStack-DsqlLambdaServiceRoleXXX'; 
GRANT USAGE ON SCHEMA testschema TO member;
GRANT SELECT, INSERT, UPDATE ON ALL TABLES IN SCHEMA testschema TO member;
exit